from flask import request, jsonify
import datetime
from database.db import get_db
import os


def count_files(path):
    count = 0
    for root, dirs, files in os.walk(path):
        for file in files:
            count += 1
    return count


def configure_dataset(app):
    # 获取所有样本集
    @app.route('/dataset', methods=['GET', 'POST', 'DELETE'])
    def dataset():
        db = get_db()
        if (request.method == 'GET'):  # 获取样本集
            original = request.args.get('original')
            datasets = db.execute('SELECT * FROM dataset').fetchall()
            data = {'dataList': []}
            if original is not None:
                original = int(original)
                filtered_datasets = [
                    dataset for dataset in datasets if
                    dataset['original'] == original and dataset['deleted'] == 0
                ]
            else:
                filtered_datasets = [
                    dataset for dataset in datasets if dataset['deleted'] == 0
                ]
            for item in filtered_datasets:
                obj = {
                    'id': item[0],
                    'original': item[1],
                    'name': item[2],
                    'path': item[3],
                    'dataBalance': item[4],
                    'classBalance': item[5],
                    'textureBalance': item[6],
                    'deleted': item[7],
                    'createTime': item[8],
                    'updateTime': item[9],
                    'deleteTime': item[10]
                }
                data['dataList'].append(obj)
            # sorted_datasets = sorted(filtered_datasets, key=lambda x: x['updateTime'], reverse=True)
            # response = {"dataList": sorted_datasets}
            return jsonify(data)
        elif request.method == 'POST':  # 插入一个样本集

            def analyze_dataset(path):
                # 在这里实现样本集分析的逻辑，计算数据均衡性、类别均衡性、纹理特征均衡性等指标
                # 返回一个包含分析结果的字典，例如：
                analysis_result = {
                    "dataBalance": 0.8,
                    "classBalance": 0.9,
                    "textureBalance": 0.7
                }
                return analysis_result

            data = request.get_json()
            path = data.get('path')
            name = data.get('name')
            if not path or not name:
                response = {"code": 1, "msg": "缺少必要参数"}
                return jsonify(response)
            try:
                analysis_result = analyze_dataset(path)
                dataset = {
                    "name": name,
                    "dataBalance": analysis_result.get("dataBalance"),
                    "classBalance": analysis_result.get("classBalance"),
                    "textureBalance": analysis_result.get("textureBalance"),
                }
                db.execute(
                    '''INSERT INTO dataset (name, path, data_Balance, class_Balance, texture_Balance)
                                VALUES (?, ?, ?, ?, ?)''',
                    (dataset["name"], path, dataset["dataBalance"],
                     dataset["classBalance"], dataset["textureBalance"]))
                db.commit()
                response = {"code": 0, "msg": "导入成功"}
                return jsonify(response)
            except Exception as e:
                response = {"code": 1, "msg": str(e)}
                return jsonify(response)
        else:
            dataset_id = request.json['id']
            db = get_db()
            # 查询样本集信息
            db.execute("SELECT * FROM dataset WHERE id=?",
                       (dataset_id, )).fetchone()
            dataset = db
            if not dataset:
                return jsonify({
                    'code': 1,
                    'msg': f"样本集 ID {dataset_id} 不存在"
                }), 400
            # 更新数据表记录
            delete_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
            db.execute(
                '''UPDATE dataset SET deleted=?, delete_Time=?
                                WHERE id=?''', (1, delete_time, dataset_id))

            db.commit()
            db.close()
            return jsonify({'code': 0, 'msg': '软删除成功'}), 200

    @app.route('/dataset/analyze', methods=['PUT'])
    def analyze_dataset():
        # 解析请求参数
        dataset_id = request.json['id']
        # 连接数据库
        db = get_db()
        # 查询样本集信息
        db.execute("SELECT * FROM dataset WHERE id=?",
                   (dataset_id, )).fetchone()
        dataset = db
        if not dataset:
            return jsonify({'code': 1, 'msg': f"样本集 ID {dataset_id} 不存在"}), 400
        # TODO: 使用之前导入的数据重新分析样本集
        # ...
        # 更新数据表记录
        data_balance = 0.5  # 假设这里是重新计算后的数据平衡值
        class_balance = 0.6  # 假设这里是重新计算后的类别平衡值
        texture_balance = 0.7  # 假设这里是重新计算后的纹理平衡值
        update_time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        db.execute(
            '''UPDATE dataset SET data_Balance=?, class_Balance=?, texture_Balance=?, update_Time=?
                        WHERE id=?''',
            (data_balance, class_balance, texture_balance, update_time,
             dataset_id))
        db.commit()
        db.close()
        return jsonify({'code': 0, 'msg': '更新成功'}), 200
